//
// Created by mio on 25-5-12.
//

#include <QDebug>
#include <QGuiApplication>
#include <QScreen>

#include "GDICapture.h"

extern "C"
{
#include <libavcodec/avcodec.h>
#include <libavdevice/avdevice.h>
#include <libavformat/avformat.h>
#include <libswscale/swscale.h>
}

GDICapture::GDICapture() : video_index_(-1), codec_context_(nullptr),
    input_format_(nullptr), format_context_(nullptr), av_frame_(nullptr), av_packet_(nullptr)
    , frame_size_(0), rgba_frame_(nullptr)
{
    // 注册设备, 使用GDI采集
    avdevice_register_all();
    // 创建frame 和 packet
    av_frame_ = std::shared_ptr<AVFrame>(av_frame_alloc(), [](AVFrame* ptr) { av_frame_free(&ptr); });
    av_packet_ = std::shared_ptr<AVPacket>(av_packet_alloc(), [](AVPacket* ptr) { av_packet_free(&ptr); });

}

GDICapture::~GDICapture()
{
    Close();
}

bool GDICapture::Init(qint32 display_index, const QRect& capture_region)
{
    if(is_initialized_)
    {
        return true;
    }

    QScreen* screen = QGuiApplication::primaryScreen();
    QRect screenGeometry = screen->geometry();
    qreal dpr = screen->devicePixelRatio();
    QSize physicalSize = QSize(screenGeometry.width() * dpr, screenGeometry.height() * dpr);
    int width = physicalSize.width();
    int height = physicalSize.height();
    QString size = QString::number(width) + "x" + QString::number(height);

    AVDictionary *options = nullptr;
    //设置属性
    av_dict_set_int(&options,"framerate",target_framerate_,AV_DICT_MATCH_CASE);//设置采集帧率
    av_dict_set_int(&options,"draw_mouse",1,AV_DICT_MATCH_CASE); //绘制鼠标
    av_dict_set_int(&options,"offset_x",0,AV_DICT_MATCH_CASE);
    av_dict_set_int(&options,"offset_y",0,AV_DICT_MATCH_CASE);
    av_dict_set(&options,"video_size",size.toStdString().c_str(),1);//屏幕分辨率

    //创建输入format
    input_format_ = const_cast<AVInputFormat*>(av_find_input_format("gdigrab"));//gdi
    if(!input_format_)
    {
        qDebug() << "av_find_input_format failed";
        return false;
    }

    //创建格式上下文
    format_context_ = avformat_alloc_context();
    if(avformat_open_input(&format_context_,"desktop",input_format_,&options) != 0)
    {
        qDebug() << "avformat_open_input failed";
        return false;
    }

    //查询流信息
    if(avformat_find_stream_info(format_context_,nullptr) < 0)
    {
        qDebug() << "avformat_find_stream_info failed";
        return false;
    }

    //找到视频流，来采集视频
    int video_index = -1;
    for(uint32_t i = 0;i<format_context_->nb_streams;i++)
    {
        if(format_context_->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO)
        {
            video_index = i;
        }
    }

    if(video_index == -1)
    {
        //没有视频流
        return false;
    }

    //创建解码器
    AVCodec* codec = const_cast<AVCodec*>(avcodec_find_decoder(format_context_->streams[video_index]->codecpar->codec_id));
    if(!codec)
    {
        return false;
    }

    //创建解码器上下文
    codec_context_ = avcodec_alloc_context3(codec);
    if(!codec_context_)
    {
        return false;
    }

    //我们需要去复制解码器上下文
    avcodec_parameters_to_context(codec_context_,format_context_->streams[video_index]->codecpar);
    //打开解码器
    if(avcodec_open2(codec_context_,codec,nullptr) != 0)
    {
        return false;
    }

    //初始化成功
    video_index_ = video_index;
    is_initialized_ = true;
    //启动线程来去捕获视频流
    this->start(); //因为我们继承QThread，所有通过start可以去执行run
    return true;
}

bool GDICapture::Close()
{
    if (is_initialized_) StopCapture();
    if (input_format_) {
        avcodec_free_context(&codec_context_);
        input_format_ = nullptr;
    }
    if (format_context_) {
        avformat_close_input(&format_context_);
        format_context_ = nullptr;
    }
    if (codec_context_) {
        avcodec_free_context(&codec_context_);
        codec_context_ = nullptr;
    }
    is_initialized_ = false;
    video_index_ = -1;
    stop_flag_ = true;
    return true;
}

bool GDICapture::CaptureFrame(AVFramePtr& bgra_frame)
{

    std::lock_guard<std::mutex> lock(mutex_);
    if (stop_flag_)
    {
        bgra_frame = nullptr;
    }

    if(rgba_frame_ == nullptr || frame_size_ == 0)
    {
        bgra_frame.reset();
        return false;
    }

    // 创建一个新的AVFrame用于输出 (确保线程安全，不直接暴露内部帧的原始指针)
    // 使用AVFramePtr进行自动管理
    AVFramePtr temp_frame_for_output(av_frame_alloc(), [](AVFrame* ptr) { if(ptr) av_frame_free(&ptr); });
    if (!temp_frame_for_output) {
        qDebug() << "GDICapture::CaptureFrame: Failed to allocate AVFrame for output.";
        bgra_frame.reset();
        return false;
    }

    temp_frame_for_output->width = this->width_;  // 使用基类成员
    temp_frame_for_output->height = this->height_; // 使用基类成员
    temp_frame_for_output->format = AV_PIX_FMT_BGRA; // GDI通常输出BGRA

    if (av_frame_get_buffer(temp_frame_for_output.get(), 0) < 0) {
        qDebug() << "GDICapture::CaptureFrame: Failed to get buffer for output AVFrame.";
        bgra_frame.reset(); // temp_frame_for_output 会自动释放
        return false;
    }

    // 从内部的 rgba_frame_ (std::shared_ptr<quint8>) 拷贝数据到新的 AVFrame
    // rgba_frame_ 中存储的是紧凑的BGRA数据
    // frame_size_ 是在Decode中根据av_frame->linesize[0] * height计算的
    // 并且Decode中的拷贝逻辑试图将数据紧凑地存入rgba_frame_

    uint8_t* src_ptr = rgba_frame_.get();
    uint8_t* dst_ptr = temp_frame_for_output->data[0];
    int dst_linesize = temp_frame_for_output->linesize[0];
    int bytes_per_pixel = 4; // BGRA

    // 确保拷贝的行字节数不超过目标帧的linesize
    // 并且源数据（rgba_frame_）是按 this->width_ * bytes_per_pixel 紧凑存储的
    size_t bytes_to_copy_per_row = (size_t)this->width_ * bytes_per_pixel;
    if (bytes_to_copy_per_row > (size_t)dst_linesize) {
        bytes_to_copy_per_row = dst_linesize; // 防止溢出，尽管对于BGRA通常相等
    }

    for (uint32_t i = 0; i < this->height_; ++i) {
        memcpy(dst_ptr + (size_t)i * dst_linesize,
               src_ptr + (size_t)i * (this->width_ * bytes_per_pixel), // 假设rgba_frame_是紧凑的
               bytes_to_copy_per_row);
    }

    // 直接移动/赋值 shared_ptr
    bgra_frame = temp_frame_for_output;

    return true;
}

void GDICapture::run() {
    if(is_initialized_ && !stop_flag_) {
        qDebug() << "GDICapture::run";
        while (!stop_flag_) {
            // std::this_thread::sleep_for(std::chrono::milliseconds(1000 / target_framerate_));
            GetOneFrame();
        }
    }
}

void GDICapture::StopCapture()
{
    if(is_initialized_)
    {
        stop_flag_ = true;
        if(this->isRunning())
        {
            this->quit();
            this->wait();
        }
        //清空这个帧数据
        std::lock_guard<std::mutex> lock(mutex_);
        rgba_frame_.reset();
        frame_size_ = 0;
        width_ = 0;
        height_ = 0;
    }
}

bool GDICapture::GetOneFrame()
{
    if (stop_flag_) return false;
    int ret = av_read_frame(format_context_, av_packet_.get());
    if (ret < 0) { return false; }
    if (av_packet_->stream_index == video_index_) {
        Decode(av_frame_.get(), av_packet_.get());
    }
    av_packet_unref(av_packet_.get());
    return true;
}

bool GDICapture::Decode(AVFrame *av_frame, AVPacket *av_packet)
{
    int ret = avcodec_send_packet(codec_context_, av_packet);
    if (ret < 0) return false;
    ret = avcodec_receive_frame(codec_context_, av_frame);
    if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) { return false; }

    // 接收成功, 需要更新帧
    // (原有的互斥锁和数据拷贝逻辑保持不变)
    { // 添加大括号以明确锁的范围，虽然FPS逻辑本身不直接访问受保护数据，但放在成功解码后是合适的
        std::lock_guard<std::mutex> lock(mutex_);
        frame_size_ = av_frame->pkt_size; // 注意：pkt_size 是压缩后的大小，对于未压缩的RGBA，可能不适用。
                                         // 通常对于未压缩的RGBA，大小是 width * height * 4
                                         // 但这里你似乎是直接从解码后的av_frame获取，我们先保持原样，
                                         // 关注点是帧计数本身。
                                         // 如果要获取原始RGBA数据大小，应该是：
                                         // frame_size_ = av_frame->width * av_frame->height * 4; (假设是RGBA)
                                         // 但你的拷贝逻辑用的是linesize[0] * height，这通常是正确的。

        rgba_frame_.reset(new uint8_t[av_frame->linesize[0] * av_frame->height], std::default_delete<uint8_t[]>());
        width_ = av_frame->width;
        height_ = av_frame->height;
        frame_size_ = av_frame->linesize[0] * av_frame->height; // 更新 frame_size_ 为实际拷贝的数据大小

        // 将frame数据拷贝出去
         // 保持你原来的拷贝逻辑，如果它对你的H264Encoder是有效的：
         // for(uint32_t i = 0; i < height_; i++) {
         //    memcpy(rgba_frame_.get() + i * width_ * 4, // 目标假设是每行 width_ * 4
         //        av_frame->data[0] + i * av_frame->linesize[0], av_frame->linesize[0]); // 源拷贝整行（可能含padding）
         // }
         // 如果 linesize[0] > width_ * 4，上述拷贝会写满目标行，但源数据可能更多。
         // 更安全的做法是确保 linesize[0] == width_ * 4，或者逐像素拷贝/使用sws_scale转换为紧凑格式。
         // 为了不改变你原有逻辑，我们暂时保留，但这是一个潜在的注意点。
         // 假设你解码出来的是 AV_PIX_FMT_BGRA (因为gdigrab通常输出这个，并且H264Encoder也接收这个)
         // 那么 linesize[0] 应该等于 width_ * 4
         if (av_frame->linesize[0] == width_ * 4) {
             memcpy(rgba_frame_.get(), av_frame->data[0], frame_size_);
         } else {
             // 如果 linesize[0] 不等于 width_ * 4，说明有 padding，需要逐行拷贝
             for (uint32_t i = 0; i < height_; i++) {
                 memcpy(rgba_frame_.get() + i * width_ * 4,
                        av_frame->data[0] + i * av_frame->linesize[0],
                        width_ * 4);
             }
         }

    } // 解锁 mutex_

    // --- FPS 统计逻辑 ---
    LogFPS();

    av_frame_unref(av_frame);
    return true;
}
